import warnings
from os import path, listdir
from copy import deepcopy
from matplotlib.pyplot import figure, draw, show, ion
from figures.matplotlib.defaults import Defaults
from numpy import zeros, array, mean, std
from h5py import File
from measurement import Measurement
from pdb import set_trace

class Experiment(dict):
    """
    Container for a set of measurements and methods 
    to analyze them.
    """
    def __init__(self):
        self.myplot = Defaults()

    @classmethod
    def from_folder(cls, folder):
        """ 
        Load the measurement files from a specified folder
        folder - string of path name
        """
        exp = cls()
        for f in listdir(folder):
            fl_nm = path.join(folder,f)
            try:
                m = Measurement.from_file(fl_nm)
            except:
                warnings.warn("Could not load file %s"%fl_nm)
            else:
                exp.add_measurement(f, m)
        return exp

    @classmethod
    def from_hdf5(cls, fname, hdf5_path):
        """
        Load measurement data from a hdf5 file.
        fname - filename of an hdf5 file
        hdf5_path - root path that contains the measurement
                    groups
        """
        exp = cls()
        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'
        
        with File(fname, mode='a') as f_h5py:
            for g in f_h5py[hdf5_path].items():
                name = str(g[0])
                exp.update({name : Measurement()})
                full_path = hdf5_path + name
                m = Measurement.from_hdf5(f_h5py, full_path)
                exp.add_measurement(name, m)
        return exp

    def to_hdf5(self, fname, hdf5_path):
        """
        Save the experiment to an hdf5 file.
        Each measurement will be saved in a sub_group
        of the hdf5_path.
        fname - filename of an hdf5 file
        hdf5_path - root path in the hdf5 file to create subgroups for
                    each experiment
        """
        with File(fname, mode='a') as f_h5py:
            for k, v in self.iteritems():
                v.to_hdf5(f_h5py, hdf5_path)
        
    @classmethod
    def from_array(cls, exparray):
        exp = cls()
        for marray in exparray:
            marray = array(marray, dtype=exparray.dtype)
            m = Measurement.from_array(marray)
            exp.add_measurement(m.name, m)
        return exp

    def to_array(self):
        """
        Create a numpy array of the measurement voltage
        """
        meas = self[self.keys()[0]]
        starray = zeros(len(self), dtype=meas.to_array().dtype)
        for i, (k, v) in enumerate(sorted(self.items())):
            starray[i] = v.to_array()
        return starray

    def add_measurement(self, name, meas):
        """
        Add a measurement to the experiment collection
        name - string key value for the added measuremen
        meas - a measurement object
        """
        self.update({name : meas})

    def subset(self, match):
        """
        Return a subeset of experiments matching the key
        match - string regular expression to match
        return Experiment object
        """
        subset = deepcopy(self)
        for k in self.iterkeys():
            if k.find(match) == -1:
                del subset[k]
        return subset

    def join(self, exp):
        """
        Add all the measurements from another experiment (exp).
        exp - Experiment object
        """
        for k, v in exp.iteritems():
            self.add_measurement(k, v)

    def shift_time(self,time):
        """
        Shift the impact voltage curve to start at the time origin (t=0)
        time - [nx1] numpy array of time
        return [nx1] time numpy array
        # TODO check
        """
        t = time - time[0]
        return t

    def scale_signal(self,signal,calibration):
        """
        Scale the signal to put into units of N
        signal - [nx1] numpy array of force related signal (lbf?)
        calibration - scalar to scale voltage to Newtons
        return [nx1] numpy array of force (N)
        # TODO check
        """
        f = signal * calibration
        return f

    def align_curves(self,x1,y1,x2,y2,frac=2./3.):
        """
        Align two curves (f1 and f2) by shifting the time axis
        x1 - [nx1] numpy array independent variable
        y1 - [nx1] numpy array dependent variable
        x2 - [nx1] numpy array independent vairable
        y2 - [nx1] numpy array dependent variable
        frac - scalar fraction of curve to use to calculate correlation (0<frac<1)
        return scalar shift y2(x2+shift) ~ y1(x1)
        # TODO check
        """
        pts = np.min([y1.shape[0],y2.shape[0]])
        if y1.shape[0]>pts:
            spline = splrep(x1,y1)
            x1 = np.linspace(x1[0],x1[-1],pts)
            y1 = splev(x1,spline)
        elif y2.shape[0]>pts:
            spline = splrep(x2,y2)
            x2 = np.linspace(x2[0],x2[-1],pts)
            y2 = splev(x2,spline)
        window_size = np.int(pts*frac)
        remainder = pts - window_size
        window = np.arange(window_size) + remainder/2
        corr = np.zeros((window_size))
        shift = np.zeros((window_size))
        i = 0
        for delta in range(-remainder/2+1,remainder/2-1):
            corr[i] = np.corrcoef(y1[window],y2[window+delta])[0,1]
            shift[i] = x1[pts/2] - x2[pts/2+delta]
            i = i + 1
        optimal_shift = shift[np.argmax(corr)]
        return optimal_shift

    def center_peaks(self,data):
        """
        Center the peaks of multiple curves in data that may not have a common x-axis
        data - [mx1] list of [nx2] numpy arrays (x,y)
        return [mx1] list of [nx2] numpy arrays (x,y)
        """
        t_max = []
        for d in data:
            t_max.append(d[0,np.argmax(d[1])])
        max_t_max = np.max(t_max)
        for i in range(len(data)):
            data[i][0] = data[i][0] + max_t_max - t_max[i]
        return data

    def plot_data(self, expression=None):
        """
        Plot the data in the experiment that have
        key names that contain expression
        expression - string that selects particular files
        """
        if expression is not None:
            subset = []
            for k, v in self.iteritems():
                if k.find(expression) != -1:
                    subset.append(k)
        else:
            subset = self.keys()

        subset.sort()

        fig = figure()
        ax = fig.add_subplot(111)
        for k in subset:
            ax.plot(self[k].raw_data[0] * 1e6, self[k].raw_data[1], label=k)
        ax.set_xlabel(r'time ($\mu$s)')
        ax.set_ylabel(r'signal (mV)')
        ax.legend()
        draw()
        show()
        return fig

    def average(self):
        """
        Average the experiments. Average the voltage signal, and estimated
        parameters. Return a measurement object where the voltage signal is the
        average of the data.
        """
        data = self.to_array()

        voltage = mean(data['voltage'], axis=0)
        std_v = std(data['voltage'], axis=0)

        try:
            Ep = mean(data['parameters']['Ep'])
            vp = mean(data['parameters']['vp'])
            COR = mean(data['parameters']['COR'])
            R2 = mean(data['parameters']['R2'])
            std_Ep = std(data['parameters']['Ep'])
            std_vp = std(data['parameters']['vp'])
            std_COR = std(data['parameters']['COR'])
            std_R2 = std(data['parameters']['R2'])
        except:
            Ep = None
            vp = None
            COR = None
            R2 = None

        example = data[data.keys()[0]]
        meas = Measurement() 
        meas.raw_data[0] = example.raw_data[0]
        meas.raw_data[1] = voltage
        meas.set_
        meas.other_data.update({'volt std dev':std_dev})
        meas.estimates['Ep'] = Ep
        meas.estimates['vp'] = vp
        meas.estimates['COR'] = COR
        meas.estimates['R2'] = R2
        return meas
